#!/usr/bin/env python
# -*- coding: utf-8 -*-

# =======================================
# file2imdb v0.1b - Copyright 2012
# Writted by muttley
# Get last version from code.google.com/p/unuseful-code/
# Thanks to IMDbPy (http://imdbpy.sourceforge.net/) used in this script 
# =======================================

import sys
import re
import os
from optparse import OptionParser


# GLOBAL VAR 
# default regex
g_cleanDate = "(.*[^ _\,\.\(\)\[\]\-])[ _\.\(\)\[\]\-]+(19[0-9][0-9]|20[0-1][0-9])([ _\,\.\(\)\[\]\-]|[^0-9]$)"
g_cleanFileName = "[ _\,\.\(\)\[\]\-](ac3|dts|custom|dc|remastered|divx|divx5|dsr|dsrip|dutch|dvd|dvd5|dvd9|dvdrip|dvdscr|dvdscreener|screener|dvdivx|cam|fragment|fs|hdtv|hdrip|hdtvrip|internal|limited|multisubs|ntsc|ogg|ogm|pal|pdtv|proper|repack|rerip|retail|r3|r5|bd5|se|svcd|swedish|german|read.nfo|nfofix|unrated|extended|ws|telesync|ts|telecine|tc|brrip|bdrip|480p|480i|576p|576i|720p|720i|1080p|1080i|3d|hrhd|hrhdtv|hddvd|bluray|x264|h264|x-264|xvid|xvidvd|xxx|www.www|cd[1-9]|\[.*\])([ _\,\.\(\)\[\]\-]|$)"
g_serieMatch = ( '\[[Ss]([0-9]+)\]_\[[Ee]([0-9]+)([^\\/]*)',                 #<!-- foo_[s01]_[e01] -->
                 '[\._ \-]([0-9]+)x([0-9]+)([^\\/]*)',                       #<!-- foo.1x09 -->
                 '[\._ \-][Ss]([0-9]+)[\.\-]?[Ee]([0-9]+)([^\\/]*)',         #<!-- foo s01e01, foo.s01.e01, foo.s01-e01 -->
                 '[\._ \-]([0-9]+)([0-9][0-9])([\._ \-][^\\/]*)',            #<!-- foo.103 -->
                 '[\._ \-]p(?:ar)?t[._ -]()([ivxlcdm]+)([\._ \-][^\\/]*)' )  #<!-- Pt.I, Part XIV -->

# default user-agent
g_mobileUserAgent = "Mozilla/5.0 (Linux; U; Android 2.2; en-us; Nexus One Build/FRF91) AppleWebKit/533.1 (KHTML, like Gecko) Version/4.0 Mobile Safari/533.1"
# default imdb url
g_imdbUrl = "http://www.imdb.com/title/tt%s/"
g_verbosity = 0

def main ():
    
    global g_verbosity, g_mobileUserAgent
    
    # Help message
    usage       = "usage: %prog [-u USERAGENT] [-v] [-s] [-b] [-p | -l | -t FILETEMPLATE] </path/to/movieFileName.ext> | <movieTitle>\n\n"\
                  "   es: %prog /path/to/movieFileName.mkv\n"\
                  "       %prog -l \"movie title (year)\"\n"\
                  "       %prog -s -t /path/to/template/file /path/to/fileSerie.mkv\n"\
                  "       %prog -p -s /path/to/fileSerie.mkv\n"\
                  "       %prog -vv -u \"Mozilla/5.0\" /path/to/movieFileName.mkv"
    
    description = "Search films or tv series on imdb and print source page, link "\
                  "or some information formatted in a template saved in a file."
    version     = " %prog 0.1b"    # --version to print
    epilog      = " Copyright 2012"
    
    # Program parsed option
    
    parser = OptionParser(usage=usage, description=description, epilog=epilog, version=version)
        
    parser.add_option("-v", "--verbose", action="count", dest="verbosity", default=g_verbosity,
                      help="Enable verbosity. Repeat for increase.")
    
    parser.add_option("-s", "--serie", action="store_true", dest="serie", default=False,
                     help="Force to search in tv series. Use only with file.\n Get parent folder name for serie title.")
    
    parser.add_option("-u", "--user-agent", type="string", dest="userAgent", default=g_mobileUserAgent,
                     help="Override default user-agent (Android Nexus one): '%s' " % g_mobileUserAgent)
    
    parser.add_option("-p", "--view-placeholder", action="store_true", dest="viewPlaceholder", default=False,
                     help="View list of possible placeholders usable in template file (es. %(rating)s).\n Add always this option in a normal search for view utilizable placeholders for particular movie.")
    
    parser.add_option("-t", "--file-template", type="string", dest="fileTemplate", default=None,
                     help="Template file with placeholders (es. %(rating)s).")
    
    parser.add_option("-l", "--print-link", action="store_true", dest="printLink", default=False,
                     help="Simple print link of finded page on imdb.")
    
    parser.add_option("-b", "--only-body", action="store_true", dest="onlyBody", default=False,
                     help="Retrive only <body> of html page.")
    
    (options, args) = parser.parse_args()
    
    # Manage params and error
    if len(args) != 1:
        myPrint( clr( 'ERROR: ', 4) + 'Select a <movieFileName.ext> or <movieTitle>!\n\n', STR_STDERR )
        parser.print_help()
        sys.exit(1)
        
    if( options.serie is True and os.access( args[0], os.F_OK ) is False):
        myPrint( clr( 'ERROR: ', 4) + 'Option --serie can be used only on file (get serie name from paren directory name)\n', STR_STDERR )
        sys.exit(1)
    
    # Init global variable
    g_mobileUserAgent = options.userAgent
    g_verbosity = options.verbosity
    movie = parentDirName = None
    movieTitleAndYear = fileName = args[0]
    
    # is a file
    if( os.access( fileName, os.F_OK ) ):
        # split fileName and path (/path/to/file.ext => file.ext)
        path, movieTitleAndYear = os.path.split( fileName )
        myPrint( clr( "INFO.02: ", 1 ) + "%s is a file...\n" % clr( movieTitleAndYear, 2 ), STR_STDOUT, 2, g_verbosity)
        # split fileName and ext
        movieTitleAndYear, strFileExt = os.path.splitext( movieTitleAndYear )
    
    # --serie option is setted
    if( options.serie is True):
        parentDirName = os.path.basename( os.path.split( os.path.abspath( fileName ) )[0] )
        myPrint( clr( "INFO.01: ", 1 ) + "Use parent directory %s as name serie\n" % clr( parentDirName, 2 ), STR_STDOUT, 1, g_verbosity)
        
        movieTitle, year = cleanFileName( parentDirName )
        season, episode = matchSeasonAndEpisodie( movieTitleAndYear )
       
        # matched season and episode
        if( season ):
            myPrint( clr( "INFO.01: ", 1 ) + "Matched season %s and episode %s from filename %s\n" % ( clr( season, 2 ), clr( episode, 2 ), clr( movieTitleAndYear, 2 ) ), STR_STDOUT, 1, g_verbosity)
            
            if( year is not None ):
                movieTitleAndYear = movieTitle + ' (' + year + ')'
        
            myPrint( clr( "INFO.01: ", 1 ) + "Searching %s on imdb (s%se%s)...\n" % ( clr( movieTitleAndYear, 2 ), season, episode ), STR_STDOUT, 1, g_verbosity)
        else:
            myPrint( clr( "INFO.01: ", 1 ) + "Searching %s on imdb (without season and episode not mached in filename)...\n" % clr( movieTitleAndYear, 2 ), STR_STDOUT, 1, g_verbosity)
            
        # search more info (RTN_BASE_CONTENT) with --view-placeholder or --file-template option setted
        movie = searchSerie( movieTitle, year, season, episode, (RTN_MIN_CONTENT, RTN_BASE_CONTENT)[ options.fileTemplate is not None or options.viewPlaceholder ] )
    else:
        myPrint( clr( "INFO.01: ", 1 ) + "Clearing (file) name...\n", STR_STDOUT, 1, g_verbosity)
        movieTitle, year = cleanFileName( movieTitleAndYear )
        
        if( year is not None ):
            movieTitleAndYear = movieTitle + ' (' + year + ')'
        
        myPrint( clr( "INFO.01: ", 1 ) + "Searching  %s on imdb...\n" % clr( movieTitleAndYear, 2 ), STR_STDOUT, 1, g_verbosity)
        
        # search more info (RTN_BASE_CONTENT) with --view-placeholder or --file-template option setted
        movie = searchMovie( movieTitle, year, (RTN_MIN_CONTENT, RTN_BASE_CONTENT)[ options.fileTemplate is not None or options.viewPlaceholder ] )
    
    if( not movie ):
        myPrint( clr( "INFO.00: ", 1 ) + "No result for: %s\n" % movieTitleAndYear, STR_STDOUT)
        sys.exit(1)
    
    else:
        # --view-placeholder option is setted
        if( options.viewPlaceholder ): 
            myPrint( clr( "INFO.00: ", 1 ) + "List of utilizable keys in template: %s\n" % clr( ', '.join( movie.keys() ), 2 ), STR_STDOUT)
        
        # --file-template option is setted (read file template and print with substitute placeholder)
        elif( options.fileTemplate is not None ):
            try:
                f = open(options.fileTemplate, 'r')
                print f.read() % movie
                f.close()
                
                myPrint( clr( "INFO.02: ", 1 ) + "Use option --view-placeholder for find utilizable placeholders.\n", STR_STDOUT, 2, g_verbosity)
            except IOError as (errno, strerror):
                myPrint( clr( 'ERROR: ', 4) + 'File template: %s (%d)\n' % ( strerror, errno ), STR_STDERR )
        
        # --print-link option is setted
        elif( options.printLink ):
            imdbUrl = g_imdbUrl % movie.movieID
            myPrint( imdbUrl + '\n', STR_STDOUT )
            
        # simple a wget with selected user-agent :)
        else:
            import urllib2
            imdbUrl = g_imdbUrl % movie.movieID
            
            myPrint( clr( "INFO.01: ", 1 ) + "Like wget %s\n" % imdbUrl, STR_STDOUT, 1, g_verbosity)
            req = urllib2.Request(url=imdbUrl, headers={'User-agent': g_mobileUserAgent})
            
            f = urllib2.urlopen(req)
            htmlPage = f.read()
            
            # retrive only sub-part of page source
            if( options.onlyBody is True):
               reSubPage = re.compile( "<body[^>]*>(.*?)</body>", re.IGNORECASE | re.DOTALL )
                mtSubPage = reSubPage.search( htmlPage )
                htmlPage = mtSubPage.group(1)
            
            print htmlPage
            
            f.close()

    sys.exit(0)

# Clear strMovieTitleAndYear from uneed tag and character
# @strMovieTitleAndYear: string - Filename ov movie film (with year if possible)
# @return: string, string - filename cleaned and year
def cleanFileName( strMovieTitleAndYear ):
    strYear = None
    
    reYear = re.compile( g_cleanDate, re.IGNORECASE )
    mtYear = reYear.search( strMovieTitleAndYear )
    
    if( mtYear is not None ):
        strMovieTitleAndYear = mtYear.group( 1 )
        strYear = mtYear.group( 2 )
    
    reTag = re.compile( g_cleanFileName, re.IGNORECASE )
    mtTag = reTag.search( strMovieTitleAndYear )
    
    if( mtTag is not None ):
        strMovieTitleAndYear = strMovieTitleAndYear[:mtTag.start( 1 ) - 1]
    
    # not find some space in title => replace [.|_] with space
    if( strMovieTitleAndYear.find(' ') == -1 ):
        strMovieTitleAndYear = strMovieTitleAndYear.replace( '.', ' ' )
        strMovieTitleAndYear = strMovieTitleAndYear.replace( '_', ' ' )
    
    strMovieTitleAndYear = strMovieTitleAndYear.strip( )
    
#    if( strYear is not None ):
#        strMovieTitleAndYear = '"' + strMovieTitleAndYear + '" (' + strYear + ')';
    
    return strMovieTitleAndYear, strYear

# Match and retunr season and episodie from filename
# @strFileName: string - filename
# @return: string, string - season, episodie
def matchSeasonAndEpisodie( strFileName ):
    strYear = None
    
    for strEpSeason in g_serieMatch:
        
        reEpSeason = re.compile( strEpSeason, re.IGNORECASE )
        mtEpSeason = reEpSeason.search( strFileName )
    
        if( mtEpSeason is not None ):
            seasons = mtEpSeason.group( 1 )
            episodie = mtEpSeason.group( 2 )
            return seasons, episodie
    
    return None, None

RTN_MIN_CONTENT = 0
RTN_BASE_CONTENT = 1
# Search movie in imdb and return movie object
# @title: string - movie title
# @year: string - movie year
# @flag: int - RTN_MIN_CONTENT, RTN_BASE_CONTENT less or more movie info
# @serie: bool - if true search in series imdb db (i think search also in film db, but if false don't search series)
# @return: object - movie object finded or 0 if nothing is finded
def searchMovie( title, year = None, flag = RTN_MIN_CONTENT, serie = False ):
    i = imdb.IMDb()

    in_encoding = sys.stdin.encoding or sys.getdefaultencoding()
    out_encoding = sys.stdout.encoding or sys.getdefaultencoding()
    
    title = unicode(title, in_encoding, 'replace')
    
    if( year is not None ):
        year = unicode(year, in_encoding, 'replace')
        if( serie is False ):
            title = title + ' (' + year + ')';
        else: # imdb compatible search for series
            title = '"' + title + '" (' + year + ')';
    
    try:
        # Do the search, and get the results (a list of Movie objects).
        results = i.search_movie( title )
    except imdb.IMDbError, e:
        myPrint( clr( 'ERROR: ', 4) + 'Probably you\'re not connected to Internet.  Complete error report: %s\n' % e, STR_STDERR )
        sys.exit(3)
    
    if not results:
        return 0
    
    movieMatched = None
    
    if( g_verbosity >= 3 or serie is True):
        myPrint( clr( "INFO.03: ", 1 ) + 'Find %s result%s for %s:\n' % ( len(results), ('', 's')[len(results) != 1], clr( title, 2 ) ), STR_STDOUT, 3, g_verbosity)
        # Print the long imdb title for every movie.
        for movie in results:
            selectedResult = ''
            
            # find selected for movies or series
            if( serie is True and selectedResult != ''): 
                selectedResult = ('', ' <-- selected')[ movie['kind'] == 'tv series' ]
            else: 
                selectedResult = ['', ' <-- selected'][ results[0] is movie ]
            
            myPrint( clr( "INFO.03: ", 1 ) + " %s :  %s [%s]%s\n" % ( i.get_imdbID(movie), movie['long imdb title'], movie['kind'], 
                                                                    selectedResult ), STR_STDOUT, 3, g_verbosity)
            
            # if searching on series on the first serie => save it and exit
            if( serie is True and movie['kind'] == 'tv series' and not movieMatched ): 
                movieMatched = movie
                if( g_verbosity < 3 ): break
     
    # get first result is not search serie
    if( serie is not True ): movieMatched = results[0]
    
    if( flag is RTN_BASE_CONTENT ):
        i.update( movieMatched )
    
    return movieMatched
    
##    from imdb.utils import analyze_title, build_title
#    #print analyze_title('"Doctor Who" (2005) {A Christmas Carol (#6.0)}') 
#    
#    #Doctor.Who.2005.6x01.L.astronauta.impossibile.ITA-ENG.720p.BDMux.h264-DarkSideMux.mkv
##    print 'movieID\t: imdbID : title'
##    
#
#    # This is a Movie instance.
#    movie = results[0]
#    url =  i.get_imdbURL(movie)
#    print movie.keys()
#    
#    # So far the Movie object only contains basic information like the
#    # title and the year; retrieve main information:
#    i.update(movie, 'episodes')
#    
#    print movie.summary().encode(out_encoding, 'replace')
#    print movie['kind'] 
#    
#    e =  movie['episodes'][6][1]
#    print i.get_imdbURL(e)
#    i.update(e)
#    print e.keys()
#    print e['rating']
##    print movie['canonical title']
##    print movie['long imdb title']
##    print movie[ 'long imdb canonical title']
##    print movie['smart canonical title']
##    print movie[ 'smart long imdb canonical title']
##    print movie[ 'long imdb episode title']

# Search specific episodie of a serie
# @title: string - serie title
# @year: string - serie year
# @season: int - serie season
# @episode: int - serie episode
# @flag: int - RTN_MIN_CONTENT, RTN_BASE_CONTENT less or more movie info
# @return: object - movie objet of episodie, or movie object of serie or 0 if nothing is finded
def searchSerie( title, year = None, season = None, episode = None, flag = RTN_MIN_CONTENT):
    movie = searchMovie( title, year, RTN_MIN_CONTENT, True )
    
    if( movie is not None ):
        # retrive specific page for episode
        if( season and episode ):
            i = imdb.IMDb()
            
            # get episodies
            i.update(movie, 'episodes')
           
            episode = movie['episodes'][int(season)][int(episode)]
            
            if( flag is RTN_BASE_CONTENT ):
                i.update( episode )
            
            return episode
        else:
            return movie
    
    return 0


STR_STDOUT = 0
STR_STDERR = 1
# Print on screen follow verbosity and stream type
# @text: string - text to print
# @streamType: int - STR_STDOUT or STR_STDERR
# @verbosity: int - verbosity associed to text
# @requestVerbosity: int - request verbosity
# @return: bool - return if printed
def myPrint( text, streamType = STR_STDOUT, verbosity = 1 , requestVerbosity = 1 ):
    
    if verbosity <= requestVerbosity:
        if streamType == STR_STDOUT:
            sys.stdout.write( text )
        else: 
            sys.stderr.write( text )
        return True
    else:
        return False
        
        
G_COLOR = ( '\033[95m', '\033[94m', '\033[92m', '\033[93m', '\033[91m' )
G_COLOR_ENDC = '\033[0m'
# Colorize string
# @string: string - text to colorize
# @color: int - color type, see G_COLOR tuple
# @return: string - return colored string
def clr( string, color ):
    
    return G_COLOR[color] + string + G_COLOR_ENDC

# Import the IMDbPY package.
try:
    import imdb
except ImportError:
    myPrint( clr( 'ERROR: ', 4) + 'You bad boy!  You need to install the IMDbPY package!\n\n', STR_STDERR )
    sys.exit(1)

if __name__ == "__main__":
    main()
